import torch

from eval import Evaler, check_orthogonality_and_determinant

tensor = torch.load('E:\\model\\models\llama2-7b\\ffn\\ffn\second_line_input\\1tensor0.pth', weights_only=True).to(torch.float32).to("cuda:0")
#
rt = torch.load("E:\\model\\models\\llama2-7b\\1-10-1 (1).pth", weights_only=True).to(torch.float32).to("cuda:0")

def random_orthogonal_matrix(dim, device):
    """生成一个随机的正交矩阵，且行列式值为 1，并使用 float32 来执行 QR 分解"""
    matrix = torch.randn(dim, dim, device=device)  # 使用 float32 来执行 QR 分解
    Q, _ = torch.linalg.qr(matrix)  # 使用 QR 分解生成正交矩阵
    # 修正行列式为 1
    if torch.det(Q) < 0:
        Q[:, 0] = -Q[:, 0]
    # 如果需要，转换结果回 float16

    return Q

def calculate_kurtosis(tensor: torch.Tensor, dim: int = None) -> torch.Tensor:
    """
    计算 PyTorch 张量的峰度 (Kurtosis)。

    参数:
    - tensor (torch.Tensor): 输入的张量 (Tensor)。
    - dim (int, 可选): 计算峰度的维度。如果为 None，则将张量展平为一维计算总体峰度。

    返回:
    - torch.Tensor: 张量的峰度值 (Kurtosis)。
    """


    if dim is not None:
        mean = torch.mean(tensor, dim=dim, keepdim=True)

        std = torch.std(tensor, dim=dim, keepdim=True)

        kurtosis = torch.mean(((tensor - mean) / std) ** 4, dim=dim)
    else:
        tensor = tensor.view(-1)

        max = torch.max(tensor)
        print("max",max.item())
        min = torch.min(tensor)
        print("min",min.item())

        mean = torch.mean(tensor)
        print("mean",mean.item())
        std = torch.std(tensor)
        print("std",std.item())
        t = ((tensor - mean) / std)

        max = torch.max(t)
        print("max-t",max.item())
        min = torch.min(t)
        print("min-t",min.item())


        print(t)
        t = t ** 4
        print(t)

        kurtosis = torch.mean(((tensor - mean) / std) ** 4)

    return kurtosis





def calculate_kurtosis_for_tensor(tensor: torch.Tensor) -> torch.Tensor:

    # 去掉 batch 维度，转换为 (sequence_length, hidden_size)
    tensor = tensor.squeeze(0)

    # print(tensor.shape)

    # 计算每个通道 (hidden_size 维度) 的均值和标准差
    mean = torch.mean(tensor, dim=1, keepdim=True)
    # print(mean.shape)
    std = torch.std(tensor, dim=1, keepdim=True) + 1e-8  # 避免除零错误
    # print(std.shape)
    # 标准化张量
    standardized_tensor = (tensor - mean) / std
    # print(standardized_tensor.shape)

    # 计算每个通道的峰度 (4 次方的均值)
    kurtosis = torch.mean(standardized_tensor ** 4, dim=1)
    # print(kurtosis.shape)

    kurtosis = torch.mean(kurtosis)
    print(kurtosis)

    return kurtosis


if __name__ == '__main__':
    # 创建一个随机张量 (标准正态分布)
    # tensor = torch.randn(1, 2, 11008)
    # evaler = Evaler()
    # result = check_orthogonality_and_determinant(rt)
    # print(result)
    # 计算张量的整体峰度
    rt = random_orthogonal_matrix(11008, "cuda:0")
    # print(tensor)

    kurtosis = calculate_kurtosis_for_tensor(tensor).item()

    print("张量的总体峰度 (Kurtosis):", kurtosis)
    print("------------")

    kurtosis2 = calculate_kurtosis_for_tensor(tensor @ rt)
    print("张量的总体峰度 (Kurtosis):", kurtosis2.item())
